package com.tencent.demo.mp4writer;

import static android.media.MediaCodec.CONFIGURE_FLAG_ENCODE;
import static android.media.MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
import static android.media.MediaFormat.KEY_BITRATE_MODE;
import static android.media.MediaFormat.KEY_BIT_RATE;
import static android.media.MediaFormat.KEY_COLOR_FORMAT;
import static android.media.MediaFormat.KEY_FRAME_RATE;
import static android.media.MediaFormat.KEY_I_FRAME_INTERVAL;
import static android.media.MediaFormat.KEY_MAX_INPUT_SIZE;
import static android.media.MediaFormat.MIMETYPE_VIDEO_AVC;
import static com.tencent.demo.mp4writer.MediaMP4Writer.TRACK_AUDIO;
import static com.tencent.demo.mp4writer.MediaMP4Writer.TRACK_VIDEO;

import android.graphics.ImageFormat;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.util.Log;

import com.tencent.trtc.TRTCCloudDef;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.LinkedBlockingQueue;

public class TRTCRemoteRecorder {
    private int videoFps;
    private int videoBitrate;
    private int audioBitrate;
    private String filePath;
    private int videoGop;
    private TRTCRemoteRecorderListener trtcRemoteRecorderListener;
    private int videoWidth=0,videoHeight=0;
    private int sampleRate=44100,channelCount=1;
    private MediaFormat videoFormat;
    private MediaFormat audioFormat;
    private MediaCodec vEncoder; //视频编码器
    private MediaCodec aEncoder; //音频编码器
    private MediaCodec.BufferInfo vBufferInfo = new MediaCodec.BufferInfo(),
            aBufferInfo = new MediaCodec.BufferInfo();
    private MediaMP4Writer mediaMP4Writer;
    private LinkedBlockingQueue<TrtcAudioFrameAndTimesTamp> audioQueue;
    private LinkedBlockingQueue<TrtcVideoFrameAndTimesTamp> videoQueue;

    private boolean audioEncoderInitialized = false; //编码器初始化成功
    private boolean videoEncoderInitialized = false;

    private boolean audioExit = false;//关闭循环
    private boolean videoExit = false;

    private boolean videoFirsFrameSuccess =false;//先从视频帧开始录制

    private boolean haveGetSpsInfo=false;//从关键帧开始录制

    private boolean ISPUSHVIDEO = false;//有远端视频流
    private boolean ISPUSHAUDIO = false;//有远端音频流

    private boolean isStartRecord = false;//开始录音

    private boolean isSpSpPs = false;//是否为spspps
    private byte[] sPsPpSBytes = null;//存放spspps

    private Thread videoEncoderThread;
    private Thread audioEncoderThread;

    private String TAGNAME = "TRTCRemoteRecorder";


    private final static TRTCRemoteRecorder INSTANCE = new TRTCRemoteRecorder();
    private TRTCRemoteRecorder(){}
    public static TRTCRemoteRecorder getInstance(){
        return INSTANCE;
    }

    public void setConfig(TRTCRecordConfig config){
        this.videoFps = config.getVideoFps();
        this.videoBitrate = config.getVideoBitrate();
        this.audioBitrate = config.getAudioBitrate();
        this.filePath = config.getFilePath();
        this.videoGop = config.getVideoGop();
    }

    public void setTrtcRemoteRecorderListener(TRTCRemoteRecorderListener trtcRemoteRecorderListener) {
        this.trtcRemoteRecorderListener = trtcRemoteRecorderListener;
    }
    public void startRecord(){
        if (ISPUSHVIDEO &&ISPUSHAUDIO){
            mediaMP4Writer = new MediaMP4Writer(filePath);
            try {
                mediaMP4Writer.init();
                if (trtcRemoteRecorderListener!=null) {
                    trtcRemoteRecorderListener.TRTCRemoteRecordStartSuccess("启动录制成功!");
                }
            } catch (IOException e) {
                e.printStackTrace();
                if (trtcRemoteRecorderListener!=null) {
                    trtcRemoteRecorderListener.TRTCRemoteRecordStartFailed("启动录制失败!");
                }
            }
            Log.i(TAGNAME,"startRecord");
            isStartRecord = true;
            audioExit = false;
            videoExit = false;
            initVideoEncoder();
            startEncoderVideo();
            initAudioEncode();
            startEncoderAudio();
        }else {
            Log.i(TAGNAME,"远端流还没有推送！！！");
        }
    }

    /**
     * 初始化音频编码
     */
    private void initAudioEncode() {
        try {
            initAudioEncoder(sampleRate,channelCount);
        } catch (IOException e) {
            e.printStackTrace();
        }

        if (aEncoder == null) {
            throw new RuntimeException("请初始化音频编码器");
        }

        aEncoder.start();
        Log.i(TAGNAME,"audioStart");
    }
    /**
     * 初始化音频编码器
     *
     * @param sampleRate  音频采样率
     * @param chanelCount 声道数
     * @throws IOException 创建编码器失败
     */
    private void initAudioEncoder(int sampleRate, int chanelCount) throws IOException {
        MediaCodec aencoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_AUDIO_AAC);
        audioFormat = MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_AAC,
                sampleRate, chanelCount);
//        audioFormat.setInteger(KEY_PCM_ENCODING, AudioFormat.ENCODING_PCM_16BIT);
        audioFormat.setInteger(KEY_MAX_INPUT_SIZE, 0);
        audioFormat.setInteger(KEY_BIT_RATE, audioBitrate);
        audioFormat.setInteger(MediaFormat.KEY_AAC_PROFILE, MediaCodecInfo.CodecProfileLevel.AACObjectLC);
        aencoder.configure(audioFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        audioQueue = new LinkedBlockingQueue<>();
        aEncoder = aencoder;
        audioEncoderInitialized =true;
        Log.i(TAGNAME,"audioEncInit");
    }

    private void startEncoderAudio(){
        audioEncoderThread = new Thread() {
            @Override
            public void run() {
                while (!audioExit) {
                    if (!audioQueue.isEmpty()&&audioEncoderInitialized) {
                        try {
                            TrtcAudioFrameAndTimesTamp data = audioQueue.take();
                            encodeAudioData(data.trtcAudioFrame, data.timestamp);
                        } catch (Exception e) {
                            e.printStackTrace();
                            if (trtcRemoteRecorderListener!=null) {
                                trtcRemoteRecorderListener.TRTCRemoteRecordAudioEncodeError("音频编码失败!");
                            }
                        }
                    }
                }
            }
        };
        audioEncoderThread.start();
    }

    private void startEncoderVideo(){
        videoEncoderThread = new Thread() {
            @Override
            public void run() {
                while (!videoExit) {
                    if (!videoQueue.isEmpty()&&videoEncoderInitialized) {
                        try {
                            TrtcVideoFrameAndTimesTamp data = videoQueue.take();
                            encodeVideoData(data.trtcVideoFrame, data.timestamp);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }

                    }
                }
            }
        };
        videoEncoderThread.start();
    }

    /**
     * 音频编码
     *
     * @param trtcAudioFrame
     */
    private void encodeAudioData(TRTCCloudDef.TRTCAudioFrame trtcAudioFrame, long timestamp) {
        Log.i(TAGNAME,"audio编码开始");
        byte[]data = trtcAudioFrame.data;
        ByteBuffer[] inputBuffers = aEncoder.getInputBuffers();
        ByteBuffer[] outputBuffers = aEncoder.getOutputBuffers();
        int inputBufferId = aEncoder.dequeueInputBuffer(-1);//缓冲区索引
        if (inputBufferId >= 0) {
            ByteBuffer inputBuffer = inputBuffers[inputBufferId];
            inputBuffer.clear();
            inputBuffer.put(data, 0, data.length);
//            long pts = new Date().getTime() * 1000 - presentationTimeUs;
            aEncoder.queueInputBuffer(inputBufferId, 0, data.length, timestamp, 0);//通知编码器编码
        }

        int outputBufferId = aEncoder.dequeueOutputBuffer(aBufferInfo, 0);
        int aacSize=0;
        byte[] aacByte =null;
        if (outputBufferId<0){
            if (trtcRemoteRecorderListener!=null){
                trtcRemoteRecorderListener.TRTCRemoteRecordAudioEncodeError("音频编码失败!");
            }
        }
        while (outputBufferId >= 0) {
            ByteBuffer outputBuffer = outputBuffers[outputBufferId];//输出的buff
            int outBitSize =aBufferInfo.size;
            byte[] audioByte = new byte[outBitSize];
            outputBuffer.get(audioByte);

            aacSize =aacSize+outBitSize;
            byte[] bytes = new byte[aacSize];
            if (aacByte!=null){
                System.arraycopy(aacByte,0,bytes,0,aacByte.length);
            }
            System.arraycopy(audioByte,0,bytes,bytes.length-audioByte.length,audioByte.length);
            aacByte = bytes;

            aEncoder.releaseOutputBuffer(outputBufferId, false);//清空缓冲区
            outputBufferId = aEncoder.dequeueOutputBuffer(aBufferInfo, 0);
        }
        if (aacByte!=null){
            if (mediaMP4Writer.mStarted&&videoFirsFrameSuccess){
                aBufferInfo.set(aBufferInfo.offset,aacSize,timestamp*1000,aBufferInfo.flags);
                aBufferInfo.size = aacByte.length;
                Log.i(TAGNAME,"encodeAudioData");
                mediaMP4Writer.writeData(TRACK_AUDIO,ByteBuffer.wrap(aacByte),aBufferInfo);
            }
        }
    }



    /**
     * 初始化视频编码
     */
    private void initVideoEncoder() {
        try {
            initVideoEncoder(videoWidth,videoHeight,videoFps,videoBitrate);
        } catch (IOException e) {
            e.printStackTrace();
            trtcRemoteRecorderListener.TRTCRemoteRecordVideoEncodeError("视频编码失败!");
        }
        if (vEncoder == null) {
            throw new RuntimeException("请初始化视频编码器");
        }
        vEncoder.start();
        Log.i(TAGNAME,"videoEncStart");

    }

    /**
     * 初始化视频编码器。
     *
     * @param width  视频的宽
     * @param height 视频的高
     * @throws IOException 创建编码器失败
     */
    public void initVideoEncoder(int width, int height, int fps, int bitrate) throws IOException {
        // 初始化
        MediaCodec vencoder = MediaCodec.createEncoderByType(MIMETYPE_VIDEO_AVC);
        videoFormat = MediaFormat.createVideoFormat(MIMETYPE_VIDEO_AVC,
                width, height);
        videoFormat.setInteger(KEY_BIT_RATE, bitrate);
        videoFormat.setInteger(KEY_COLOR_FORMAT, COLOR_FormatYUV420SemiPlanar);//颜色格式
        videoFormat.setInteger(KEY_FRAME_RATE, fps);
        videoFormat.setInteger(KEY_I_FRAME_INTERVAL, videoGop);//I帧将要出现的时间间隔
        videoFormat.setInteger(MediaFormat.KEY_PROFILE, MediaCodecInfo.CodecProfileLevel.AVCProfileHigh);
        videoFormat.setInteger("level", MediaCodecInfo.CodecProfileLevel.AVCLevel41); // Level 4.1

        /**此处framework层有个bug，framework内部把支持的模式写死了  不能使用isBitrateModeSupported来判断手机支持的码率模式。
         * 首选cq 不支持的手机在用vbr
        **/
        try {
            videoFormat.setInteger(KEY_BITRATE_MODE,MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_CQ);
            vencoder.configure(videoFormat, null, null, CONFIGURE_FLAG_ENCODE);
        }catch (Exception e){
            videoFormat.setInteger(KEY_BITRATE_MODE,MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_VBR);
            vencoder.configure(videoFormat, null, null, CONFIGURE_FLAG_ENCODE);
        }
//        vencoder.configure(videoFormat, null, null, CONFIGURE_FLAG_ENCODE);
        videoQueue = new LinkedBlockingQueue<>();
        vEncoder = vencoder;
        videoEncoderInitialized =true;
        Log.i(TAGNAME,"videoEncInit");
    }



    /**
     * 视频编码
     *
     * @param trtcVideoFrame
     */
    private void encodeVideoData(TRTCCloudDef.TRTCVideoFrame trtcVideoFrame, long timestamp) {
        Log.i(TAGNAME,"video编码开始");
        byte[] data = trtcVideoFrame.data;
        byte[] dstByte = new byte[calculateFrameSize(ImageFormat.NV21,trtcVideoFrame.width,trtcVideoFrame.height)];
        Yuv420Util.I420ToNV12(data,dstByte,trtcVideoFrame.width,trtcVideoFrame.height);
        ByteBuffer[] inputBuffers = vEncoder.getInputBuffers();
        ByteBuffer[] outputBuffers = vEncoder.getOutputBuffers();

        int inputBufferId = vEncoder.dequeueInputBuffer(-1);
        if (inputBufferId >= 0) {
            // fill inputBuffers[inputBufferId] with valid data
            ByteBuffer bb = inputBuffers[inputBufferId];

            bb.clear();
            bb.put(dstByte, 0, dstByte.length);
            vBufferInfo.presentationTimeUs = timestamp*1000;
            vEncoder.queueInputBuffer(inputBufferId, 0, dstByte.length, timestamp*1000, 0);
        }

        int outputBufferId = vEncoder.dequeueOutputBuffer(vBufferInfo, 0);
        if (outputBufferId ==MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
            mediaMP4Writer.setVideoTrack(vEncoder.getOutputFormat());
            mediaMP4Writer.setAudioTrack(aEncoder.getOutputFormat());
            Log.i("mediaMP4Writer","setAudioTrack");
            Log.i("mediaMP4Writer","setVideoTrack");
            mediaMP4Writer.start();
        }
        int h264Size=0;
        byte[] h264Byte =null;
        if (outputBufferId<0){
            if (trtcRemoteRecorderListener!=null){
                trtcRemoteRecorderListener.TRTCRemoteRecordVideoEncodeError("视频编码失败!");
            }
        }
        Log.i(TAGNAME,"循环前的flag："+vBufferInfo.flags);
        while (outputBufferId >= 0) {
            // outputBuffers[outputBufferId] is ready to be processed or rendered.
            ByteBuffer bb = outputBuffers[outputBufferId];
            byte[] outData = new byte[vBufferInfo.size];
            bb.get(outData);
            Log.i(TAGNAME,"循环中的flag："+vBufferInfo.flags);
            //判断是否为spspps  是的话直接写入MP4
            if (vBufferInfo.flags == MediaCodec.BUFFER_FLAG_CODEC_CONFIG){
                if (outData.length>0&&mediaMP4Writer.mStarted){
                    isSpSpPs = true;
                    vBufferInfo.set(vBufferInfo.offset,outData.length,timestamp*1000,vBufferInfo.flags);
                    videoFirsFrameSuccess = true;
                    sPsPpSBytes = outData;
                    mediaMP4Writer.writeData(TRACK_VIDEO,ByteBuffer.wrap(outData),vBufferInfo);
                    vEncoder.releaseOutputBuffer(outputBufferId, false);
                    outputBufferId = vEncoder.dequeueOutputBuffer(vBufferInfo, 0);
                    continue;
                }
            }
            //判断是否为关键帧 是的话在前面插入spspps
            if (vBufferInfo.flags == MediaCodec.BUFFER_FLAG_KEY_FRAME){
                if (sPsPpSBytes!=null&&!isSpSpPs){
                    vBufferInfo.set(vBufferInfo.offset,sPsPpSBytes.length,timestamp*1000,MediaCodec.BUFFER_FLAG_CODEC_CONFIG);
                    mediaMP4Writer.writeData(TRACK_VIDEO,ByteBuffer.wrap(sPsPpSBytes),vBufferInfo);
                    vBufferInfo.set(vBufferInfo.offset,outData.length,timestamp*1000,MediaCodec.BUFFER_FLAG_KEY_FRAME);

                }
            }
            //将缓冲区生产的数据合并为一帧
            h264Size =h264Size+outData.length;
            byte[] bytes = new byte[h264Size];
            if (h264Byte!=null){
                System.arraycopy(h264Byte,0,bytes,0,h264Byte.length);
            }
            System.arraycopy(outData,0,bytes,bytes.length-outData.length,outData.length);
            h264Byte = bytes;
            vEncoder.releaseOutputBuffer(outputBufferId, false);
            vBufferInfo.set(vBufferInfo.offset,outData.length,timestamp*1000,vBufferInfo.flags);
            outputBufferId = vEncoder.dequeueOutputBuffer(vBufferInfo, 0);
        }

        Log.i(TAGNAME,"循环后的flag："+vBufferInfo.flags);

        if (h264Byte!=null&& mediaMP4Writer.mStarted){
            vBufferInfo.set(vBufferInfo.offset,h264Byte.length,timestamp*1000,vBufferInfo.flags);
            videoFirsFrameSuccess = true;
            mediaMP4Writer.writeData(TRACK_VIDEO,ByteBuffer.wrap(h264Byte),vBufferInfo);
            isSpSpPs = false;
        }
    }

    //设值缓存区大小
    private int calculateFrameSize(int format,int width,int height) {
        return width * height * ImageFormat.getBitsPerPixel(format) / 8;
    }


    public void stopRecord(){
        if(mediaMP4Writer.mStarted){
            try {
                mediaMP4Writer.stop();
                if (trtcRemoteRecorderListener!=null) {
                    trtcRemoteRecorderListener.TRTCRemoteRecordStopSuccess("停止录制成功!");
                }
            }catch (Exception e){
                e.printStackTrace();
                if (trtcRemoteRecorderListener!=null) {
                    trtcRemoteRecorderListener.TRTCRemoteRecordStopFailed("停止录制失败!");
                }
            }
        }
        isStartRecord = false;
        Log.i(TAGNAME,"stopRecord");

        if (aEncoder!=null){
            try {
                aEncoder.stop();
                aEncoder.release();
                vEncoder.stop();
                vEncoder.release();
            }catch (Exception  e){
                e.printStackTrace();
            }

        }
        videoEncoderThread.interrupt();
        audioEncoderThread.interrupt();
        audioExit = true;
        videoExit = true;
        audioQueue.clear();
        videoQueue.clear();
    }
    public void pushVideoFrame(TRTCCloudDef.TRTCVideoFrame trtcVideoFrame,long trtcVideoCustomPTS){
        ISPUSHVIDEO = true;
        Log.i(TAGNAME,trtcVideoFrame.toString());
        videoWidth = trtcVideoFrame.width;
        videoHeight = trtcVideoFrame.height;
        if (isStartRecord){
            try {
                TrtcVideoFrameAndTimesTamp trtcVideoFrameAndTimesTamp = new TrtcVideoFrameAndTimesTamp(trtcVideoFrame, trtcVideoCustomPTS);
                videoQueue.put(trtcVideoFrameAndTimesTamp);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public void pushAudioFrame(TRTCCloudDef.TRTCAudioFrame trtcAudioFrame, long trtcAudioCustomPTS){
        ISPUSHAUDIO = true;
        sampleRate = trtcAudioFrame.sampleRate;
        channelCount = trtcAudioFrame.channel;
        if (isStartRecord){
            try {
                audioQueue.put(new TrtcAudioFrameAndTimesTamp(trtcAudioFrame,trtcAudioCustomPTS));
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public interface TRTCRemoteRecorderListener{
        void TRTCRemoteRecordStartSuccess(String message);// 启动录制成功
        void TRTCRemoteRecordStopSuccess(String message);// 停止录制成功
        void TRTCRemoteRecordStartFailed(String message);// 启动录制失败
        void TRTCRemoteRecordStopFailed(String message);// 停止录制失败
        void TRTCRemoteRecordAudioEncodeError(String message);// 音频编码失败
        void TRTCRemoteRecordVideoEncodeError(String message);// 视频编码失败
        void TRTCRemoteRecordVideoMuxerError(String message);// 合成Mp4失败
        void TRTCRemoteRecordComplete(String message);// 合成Mp4成功

    }


    class TrtcAudioFrameAndTimesTamp{
        public TrtcAudioFrameAndTimesTamp(TRTCCloudDef.TRTCAudioFrame trtcAudioFrame, long timestamp) {
            this.trtcAudioFrame = trtcAudioFrame;
            this.timestamp = timestamp;
        }

        @Override
        public String toString() {
            return "TrtcAudioFrameAndTimesTamp{" +
                    "trtcAudioFrame=" + trtcAudioFrame +
                    ", timestamp=" + timestamp +
                    '}';
        }

        public TRTCCloudDef.TRTCAudioFrame getTrtcAudioFrame() {
            return trtcAudioFrame;
        }

        public void setTrtcAudioFrame(TRTCCloudDef.TRTCAudioFrame trtcAudioFrame) {
            this.trtcAudioFrame = trtcAudioFrame;
        }

        public long getTimestamp() {
            return timestamp;
        }

        public void setTimestamp(long timestamp) {
            this.timestamp = timestamp;
        }

        TRTCCloudDef.TRTCAudioFrame trtcAudioFrame;
        long timestamp;
    }

    class TrtcVideoFrameAndTimesTamp{
        @Override
        public String toString() {
            return "TrtcVideoFrameAndTimesTamp{" +
                    "trtcVideoFrame=" + trtcVideoFrame +
                    ", timestamp=" + timestamp +
                    '}';
        }

        public TRTCCloudDef.TRTCVideoFrame getTrtcVideoFrame() {
            return trtcVideoFrame;
        }

        public void setTrtcVideoFrame(TRTCCloudDef.TRTCVideoFrame trtcVideoFrame) {
            this.trtcVideoFrame = trtcVideoFrame;
        }

        public long getTimestamp() {
            return timestamp;
        }

        public void setTimestamp(long timestamp) {
            this.timestamp = timestamp;
        }

        public TrtcVideoFrameAndTimesTamp(TRTCCloudDef.TRTCVideoFrame trtcVideoFrame, long timestamp) {
            this.trtcVideoFrame = trtcVideoFrame;
            this.timestamp = timestamp;
        }

        TRTCCloudDef.TRTCVideoFrame trtcVideoFrame;
        long timestamp;
    }
}
